package com.yunyi.handler;

import com.yunyi.entity.mysql.ChatMessageEntity;
import com.yunyi.mapper.ChatMessageMysqlMapper;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;

import java.time.LocalDateTime;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

@RequiredArgsConstructor
@Component
public class DbChatMemoryHandler implements ChatMemory {


    private final ChatMessageMysqlMapper messageMapper;


    @Override
    public void add(String conversationId, List<Message> messages) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        Assert.notNull(messages, "messages cannot be null");
        Assert.noNullElements(messages, "messages cannot contain null elements");

        List<ChatMessageEntity> entities = messages.stream()
                .map(msg -> {
                    ChatMessageEntity entity = new ChatMessageEntity();
                    entity.setConversationId(conversationId);
                    entity.setContent(msg.getText());

                    if (msg instanceof UserMessage) {
                        entity.setMessageType(ChatMessageEntity.MessageType.USER);
                    } else if (msg instanceof AssistantMessage) {
                        entity.setMessageType(ChatMessageEntity.MessageType.ASSISTANT);
                    }

                    entity.setCreatedAt(LocalDateTime.now());
                    return entity;
                })
                .collect(Collectors.toList());

        messageMapper.insertMessages(entities);
    }

    @Override
    public List<Message> get(String conversationId) {

        List<ChatMessageEntity> entities = messageMapper.findLastNMessages(conversationId);
        Collections.reverse(entities);

        return entities.stream()
                .map(entity -> {
                    switch (entity.getMessageType()) {
                        case USER:
                            return new UserMessage(entity.getContent());
                        case ASSISTANT:
                            return new AssistantMessage(entity.getContent());
                        default:
                            throw new IllegalArgumentException("未知的消息类型");
                    }
                })
                .collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        messageMapper.deleteByConversationId(conversationId);
    }
}
